#!/usr/bin/env python

import os, sys, getopt
import csv, json

import utility


# given file index,
# return correspondence file name
def generate_file_name_by_index(index):
    file_name_prefix = "part-00"
    file_name_suffix = "-of-00500.csv"

    if index < 10:
        file_name = file_name_prefix+"00"+str(index)+file_name_suffix
    elif index < 100:
        file_name = file_name_prefix+"0"+str(index)+file_name_suffix
    else:
        file_name = file_name_prefix+str(index)+file_name_suffix
    return file_name


# read task usage and write to corrospondence folder
# jobs {
# -----job_id
# -----}
def write_task_usage_to_job_folder(task_usage, output_prefix):
    job_id_to_info = {}
    for usage in task_usage:
        job_id = usage[0]
        if job_id not in job_id_to_info:
            job_id_to_info[job_id] = []
        job_id_to_info[job_id].append(usage)

    for (job_id, usage) in job_id_to_info.items():
        with open(output_prefix+job_id, "a+") as f:
            for record in usage:
                content = record[0] + " " + record[1] + " " \
                        + record[2] + " " + record[3] + "\n"
                f.write(content)


# read task events and write to corrospondence folder
# user_id {
# --------priority {
# -----------------job_id
# -----------------}
# --------}
def write_task_events_to_user_folder(task_events, hash_value_to_user_id, output_prefix):
    job_id_to_info = {}
    for event in task_events:
        (user, job_id, priority) = (event[0], event[1], event[2])
        job_id_to_info[job_id] = {
            "user":     user,
            "priority": priority
        }

    for (job_id, info) in job_id_to_info.items():
        if info["user"] not in hash_value_to_user_id:
            user_index = str(len(hash_value_to_user_id))
            user_length = len(user_index)
            if user_length == 1:
                user_index = "00"+user_index
            elif user_length == 2:
                user_index = "0"+user_index
            hash_value_to_user_id[info["user"]] = user_index
            os.system("mkdir "+output_prefix+user_index)
        
        with open(output_prefix+hash_value_to_user_id[info["user"]]+"/"+info["priority"], "a+") as f:
            f.write(job_id+"\n")


# given user_folder and job_folder
# move corrospondence job from job folder to user folder by priority
def move_corrospondence_job_to_user_folder(user_prefix, job_prefix):
    priority_levels = os.listdir(user_prefix)
    for priority_level in priority_levels:
        target_folder = user_prefix+"priority"+priority_level
        os.system("mkdir "+target_folder)
        with open(user_prefix+priority_level) as f:
            for job_id in f:
                job_id = job_id.strip('\n')
                original_file = job_prefix+str(job_id)
                if os.path.exists(original_file):
                    os.system("mv "+original_file+" "+target_folder)


# main framework that trigger the process
def process_data(begin_index, end_index):
    task_events_prefix = "task_events/"
    task_usage_prefix = "task_usage/"
    
    output_prefix = "output/"
    if not os.path.exists(output_prefix):
        os.system("mkdir "+output_prefix)

    hash_value_to_user_id = {}
    for index in range(begin_index, end_index):
        task_events_file_name = task_events_prefix+generate_file_name_by_index(index)
        task_events = utility.extract_task_events(task_events_file_name)
        write_task_events_to_user_folder(task_events, hash_value_to_user_id, output_prefix)


    job_prefix = output_prefix+"jobs/"
    if not os.path.exists(job_prefix):
        os.system("mkdir "+job_prefix)    
    for index in range(begin_index, end_index):
        task_usage_file_name = task_usage_prefix+generate_file_name_by_index(index)
        task_usage = utility.extract_task_usage(task_usage_file_name)
        write_task_usage_to_job_folder(task_usage, job_prefix)

    user_files = os.listdir(output_prefix)
    user_files.remove("jobs")
    for user_file in user_files:
        move_corrospondence_job_to_user_folder(output_prefix+user_file+"/", job_prefix)


def main(argv):
    try:
        opts, args = getopt.getopt(argv[1:], "hb:e:")
    except getopt.GetoptError as err:
        print(err)
        sys.exit(2)

    # set default value
    begin_index = 0
    end_index = 500

    for o, s in opts:
        # help information
        if o == "-h":
            print "usge   -b <begin package index>"
            print "       -e <end package index>"
            return
        elif o == "-b":
            begin_index = int(s)
        elif o == "-e":
            end_index = int(s)

    process_data(begin_index, end_index)
    

if __name__ == "__main__":
    main(sys.argv)
